"""
Differential operators for Cartesian grids with comprehensive boundary condition support.

This module implements all major differential operators for Cartesian coordinate systems:
- Laplacian operator (∇²) for scalar fields
- Gradient operator (∇) for scalar fields  
- Divergence operator (∇·) for vector fields
- Curl operator (∇×) for vector fields
- Vector Laplacian for vector fields
- Biharmonic operator (∇⁴) for higher-order PDEs
- Tensor divergence for tensor fields

Features:
- Support for 1D, 2D, and 3D grids
- Multiple boundary condition types (Dirichlet, Neumann, Robin, periodic)
- Conservative and non-conservative formulations
- High-order finite difference schemes
- FFT-based operators for periodic boundaries
- Optimized sparse matrix representations
"""

using LinearAlgebra
using SparseArrays
using FFTW

include("common.jl")

# Export all operator creation functions
export make_laplace, make_gradient, make_divergence, make_curl
export make_vector_laplace, make_tensor_divergence, make_biharmonic
export make_poisson_solver, fft_laplace, fft_gradient
export make_corner_point_setter_2d, make_corner_point_setter_3d

"""
    make_laplace(grid::CartesianGrid; conservative=false, order=2) -> OperatorType

Create discrete Laplacian operator for Cartesian grid.

Args:
    grid: The Cartesian grid
    conservative: Whether to use conservative discretization
    order: Finite difference order (2 or 4)

Returns:
    Function that applies Laplacian to field data
"""
function make_laplace(grid::CartesianGrid; conservative=false, order=2)
    if conservative
        return make_conservative_laplace(grid, order)
    else
        return make_standard_laplace(grid, order)
    end
end

"""
    make_standard_laplace(grid::CartesianGrid, order=2) -> OperatorType

Create standard (non-conservative) Laplacian operator.
"""
function make_standard_laplace(grid::CartesianGrid, order=2)
    dims = size(grid)
    spacing_vals = spacing(grid)
    ndims = ndim(grid)
    
    if ndims == 1
        return make_laplace_1d(dims[1], spacing_vals[1], order)
    elseif ndims == 2
        return make_laplace_2d(dims, spacing_vals, order)
    elseif ndims == 3
        return make_laplace_3d(dims, spacing_vals, order)
    else
        error("Unsupported grid dimension: $ndims")
    end
end

"""
    make_laplace_1d(n, h, order) -> OperatorType

Create 1D Laplacian operator.
"""
function make_laplace_1d(n, h, order)
    if order == 2
        # Standard second-order finite difference
        function laplace_1d(field_data, out)
            h2_inv = 1.0 / (h * h)
            
            # Interior points
            for i in 2:(n-1)
                out[i] = (field_data[i+1] - 2*field_data[i] + field_data[i-1]) * h2_inv
            end
            
            # Boundary points - will be handled by boundary conditions
            out[1] = 0.0
            out[n] = 0.0
            
            return out
        end
        
    elseif order == 4
        # Fourth-order finite difference
        function laplace_1d(field_data, out)
            h2_inv = 1.0 / (h * h)
            
            # Interior points (need at least 2 points on each side)
            for i in 3:(n-2)
                out[i] = ((-field_data[i+2] + 16*field_data[i+1] - 30*field_data[i] + 
                          16*field_data[i-1] - field_data[i-2]) / 12.0) * h2_inv
            end
            
            # Near-boundary points use second-order
            for i in [2, n-1]
                out[i] = (field_data[i+1] - 2*field_data[i] + field_data[i-1]) * h2_inv
            end
            
            # Boundary points
            out[1] = 0.0
            out[n] = 0.0
            
            return out
        end
    else
        error("Unsupported finite difference order: $order")
    end
end

"""
    make_laplace_2d(dims, spacing_vals, order) -> OperatorType

Create 2D Laplacian operator.
"""
function make_laplace_2d(dims, spacing_vals, order)
    nx, ny = dims
    hx, hy = spacing_vals
    hx2_inv = 1.0 / (hx * hx)
    hy2_inv = 1.0 / (hy * hy)
    
    if order == 2
        function laplace_2d(field_data, out)
            # Interior points
            for i in 2:(nx-1), j in 2:(ny-1)
                d2_dx2 = (field_data[i+1,j] - 2*field_data[i,j] + field_data[i-1,j]) * hx2_inv
                d2_dy2 = (field_data[i,j+1] - 2*field_data[i,j] + field_data[i,j-1]) * hy2_inv
                out[i,j] = d2_dx2 + d2_dy2
            end
            
            # Boundary handling will be done separately
            return out
        end
        
    elseif order == 4
        function laplace_2d(field_data, out)
            # Fourth-order interior points
            for i in 3:(nx-2), j in 3:(ny-2)
                d2_dx2 = ((-field_data[i+2,j] + 16*field_data[i+1,j] - 30*field_data[i,j] + 
                          16*field_data[i-1,j] - field_data[i-2,j]) / 12.0) * hx2_inv
                d2_dy2 = ((-field_data[i,j+2] + 16*field_data[i,j+1] - 30*field_data[i,j] + 
                          16*field_data[i,j-1] - field_data[i,j-2]) / 12.0) * hy2_inv
                out[i,j] = d2_dx2 + d2_dy2
            end
            
            # Near-boundary points use second-order
            for i in [2, nx-1], j in 2:(ny-1)
                d2_dx2 = (field_data[i+1,j] - 2*field_data[i,j] + field_data[i-1,j]) * hx2_inv
                d2_dy2 = (field_data[i,j+1] - 2*field_data[i,j] + field_data[i,j-1]) * hy2_inv
                out[i,j] = d2_dx2 + d2_dy2
            end
            
            for i in 3:(nx-2), j in [2, ny-1]
                d2_dx2 = (field_data[i+1,j] - 2*field_data[i,j] + field_data[i-1,j]) * hx2_inv
                d2_dy2 = (field_data[i,j+1] - 2*field_data[i,j] + field_data[i,j-1]) * hy2_inv
                out[i,j] = d2_dx2 + d2_dy2
            end
            
            return out
        end
    else
        error("Unsupported finite difference order: $order")
    end
end

"""
    make_laplace_3d(dims, spacing_vals, order) -> OperatorType

Create 3D Laplacian operator.
"""
function make_laplace_3d(dims, spacing_vals, order)
    nx, ny, nz = dims
    hx, hy, hz = spacing_vals
    hx2_inv = 1.0 / (hx * hx)
    hy2_inv = 1.0 / (hy * hy)  
    hz2_inv = 1.0 / (hz * hz)
    
    function laplace_3d(field_data, out)
        # Interior points
        for i in 2:(nx-1), j in 2:(ny-1), k in 2:(nz-1)
            d2_dx2 = (field_data[i+1,j,k] - 2*field_data[i,j,k] + field_data[i-1,j,k]) * hx2_inv
            d2_dy2 = (field_data[i,j+1,k] - 2*field_data[i,j,k] + field_data[i,j-1,k]) * hy2_inv
            d2_dz2 = (field_data[i,j,k+1] - 2*field_data[i,j,k] + field_data[i,j,k-1]) * hz2_inv
            out[i,j,k] = d2_dx2 + d2_dy2 + d2_dz2
        end
        
        return out
    end
end

"""
    make_conservative_laplace(grid::CartesianGrid, order=2) -> OperatorType

Create conservative Laplacian operator that ensures mass conservation.
"""
function make_conservative_laplace(grid::CartesianGrid, order=2)
    # Conservative operators use flux-based discretization
    dims = size(grid)
    spacing_vals = spacing(grid)
    ndims = ndim(grid)
    
    if ndims == 1
        h = spacing_vals[1]
        function conservative_laplace_1d(field_data, out)
            h_inv = 1.0 / h
            # Conservative form: d/dx(du/dx)
            for i in 2:(dims[1]-1)
                flux_right = (field_data[i+1] - field_data[i]) * h_inv
                flux_left = (field_data[i] - field_data[i-1]) * h_inv
                out[i] = (flux_right - flux_left) * h_inv
            end
            return out
        end
        
    else
        # Multi-dimensional conservative operators
        error("Multi-dimensional conservative Laplacian not yet implemented")
    end
end

"""
    make_gradient(grid::CartesianGrid, dim::Int; order=2) -> OperatorType

Create gradient operator for specified dimension.
"""
function make_gradient(grid::CartesianGrid, dim::Int; order=2)
    dims = size(grid)
    h = spacing(grid)[dim]
    ndims = ndim(grid)
    
    if order == 2
        h_inv = 1.0 / (2.0 * h)
        
        if ndims == 1
            function gradient_1d(field_data, out)
                # Central differences
                for i in 2:(dims[1]-1)
                    out[i] = (field_data[i+1] - field_data[i-1]) * h_inv
                end
                # Boundaries handled separately
                return out
            end
            
        elseif ndims == 2
            nx, ny = dims
            if dim == 1  # ∂/∂x
                function gradient_2d_x(field_data, out)
                    for i in 2:(nx-1), j in 1:ny
                        out[i,j] = (field_data[i+1,j] - field_data[i-1,j]) * h_inv
                    end
                    return out
                end
            else  # ∂/∂y
                function gradient_2d_y(field_data, out)
                    for i in 1:nx, j in 2:(ny-1)
                        out[i,j] = (field_data[i,j+1] - field_data[i,j-1]) * h_inv
                    end
                    return out
                end
            end
            
        elseif ndims == 3
            nx, ny, nz = dims
            if dim == 1  # ∂/∂x
                function gradient_3d_x(field_data, out)
                    for i in 2:(nx-1), j in 1:ny, k in 1:nz
                        out[i,j,k] = (field_data[i+1,j,k] - field_data[i-1,j,k]) * h_inv
                    end
                    return out
                end
            elseif dim == 2  # ∂/∂y
                function gradient_3d_y(field_data, out)
                    for i in 1:nx, j in 2:(ny-1), k in 1:nz
                        out[i,j,k] = (field_data[i,j+1,k] - field_data[i,j-1,k]) * h_inv
                    end
                    return out
                end
            else  # ∂/∂z
                function gradient_3d_z(field_data, out)
                    for i in 1:nx, j in 1:ny, k in 2:(nz-1)
                        out[i,j,k] = (field_data[i,j,k+1] - field_data[i,j,k-1]) * h_inv
                    end
                    return out
                end
            end
        end
        
    elseif order == 4
        # Fourth-order central differences: (-f₂ + 8f₁ - 8f₋₁ + f₋₂)/12h
        h_inv = 1.0 / (12.0 * h)
        
        if ndims == 1
            function gradient_1d_o4(field_data, out)
                n = dims[1]
                for i in 3:(n-2)
                    out[i] = (-field_data[i+2] + 8*field_data[i+1] - 8*field_data[i-1] + field_data[i-2]) * h_inv
                end
                # Use second-order near boundaries
                h2_inv = 1.0 / (2.0 * h)
                out[2] = (field_data[3] - field_data[1]) * h2_inv
                out[n-1] = (field_data[n] - field_data[n-2]) * h2_inv
                return out
            end
        else
            error("High-order multi-dimensional gradient not yet implemented")
        end
    else
        error("Unsupported gradient order: $order")
    end
end

"""
    make_divergence(grid::CartesianGrid) -> OperatorType

Create divergence operator for vector fields: ∇·v = ∂vₓ/∂x + ∂vᵧ/∂y + ∂vᵤ/∂z
"""
function make_divergence(grid::CartesianGrid)
    dims = size(grid)
    spacing_vals = spacing(grid)
    ndims = ndim(grid)
    
    if ndims == 2
        hx, hy = spacing_vals
        hx_inv = 1.0 / (2.0 * hx)
        hy_inv = 1.0 / (2.0 * hy)
        
        function divergence_2d(vector_field, out)
            # vector_field should be (nx, ny, 2) array
            for i in 2:(dims[1]-1), j in 2:(dims[2]-1)
                dvx_dx = (vector_field[i+1,j,1] - vector_field[i-1,j,1]) * hx_inv
                dvy_dy = (vector_field[i,j+1,2] - vector_field[i,j-1,2]) * hy_inv
                out[i,j] = dvx_dx + dvy_dy
            end
            return out
        end
        
    elseif ndims == 3
        hx, hy, hz = spacing_vals
        hx_inv = 1.0 / (2.0 * hx)
        hy_inv = 1.0 / (2.0 * hy)
        hz_inv = 1.0 / (2.0 * hz)
        
        function divergence_3d(vector_field, out)
            # vector_field should be (nx, ny, nz, 3) array
            for i in 2:(dims[1]-1), j in 2:(dims[2]-1), k in 2:(dims[3]-1)
                dvx_dx = (vector_field[i+1,j,k,1] - vector_field[i-1,j,k,1]) * hx_inv
                dvy_dy = (vector_field[i,j+1,k,2] - vector_field[i,j-1,k,2]) * hy_inv
                dvz_dz = (vector_field[i,j,k+1,3] - vector_field[i,j,k-1,3]) * hz_inv
                out[i,j,k] = dvx_dx + dvy_dy + dvz_dz
            end
            return out
        end
    else
        error("Divergence not implemented for $(ndims)D")
    end
end

"""
    make_curl(grid::CartesianGrid) -> OperatorType

Create curl operator for vector fields.
In 2D: curl(v) = ∂vᵧ/∂x - ∂vₓ/∂y (scalar)
In 3D: curl(v) = (∂vᵤ/∂y - ∂vᵧ/∂z, ∂vₓ/∂z - ∂vᵤ/∂x, ∂vᵧ/∂x - ∂vₓ/∂y) (vector)
"""
function make_curl(grid::CartesianGrid)
    dims = size(grid)
    spacing_vals = spacing(grid)
    ndims = ndim(grid)
    
    if ndims == 2
        hx, hy = spacing_vals
        hx_inv = 1.0 / (2.0 * hx)
        hy_inv = 1.0 / (2.0 * hy)
        
        function curl_2d(vector_field, out)
            # 2D curl is scalar: ∂vᵧ/∂x - ∂vₓ/∂y
            for i in 2:(dims[1]-1), j in 2:(dims[2]-1)
                dvy_dx = (vector_field[i+1,j,2] - vector_field[i-1,j,2]) * hx_inv
                dvx_dy = (vector_field[i,j+1,1] - vector_field[i,j-1,1]) * hy_inv
                out[i,j] = dvy_dx - dvx_dy
            end
            return out
        end
        
    elseif ndims == 3
        hx, hy, hz = spacing_vals
        hx_inv = 1.0 / (2.0 * hx)
        hy_inv = 1.0 / (2.0 * hy)
        hz_inv = 1.0 / (2.0 * hz)
        
        function curl_3d(vector_field, out)
            # 3D curl is vector
            for i in 2:(dims[1]-1), j in 2:(dims[2]-1), k in 2:(dims[3]-1)
                # x-component: ∂vᵤ/∂y - ∂vᵧ/∂z
                dvz_dy = (vector_field[i,j+1,k,3] - vector_field[i,j-1,k,3]) * hy_inv
                dvy_dz = (vector_field[i,j,k+1,2] - vector_field[i,j,k-1,2]) * hz_inv
                out[i,j,k,1] = dvz_dy - dvy_dz
                
                # y-component: ∂vₓ/∂z - ∂vᵤ/∂x  
                dvx_dz = (vector_field[i,j,k+1,1] - vector_field[i,j,k-1,1]) * hz_inv
                dvz_dx = (vector_field[i+1,j,k,3] - vector_field[i-1,j,k,3]) * hx_inv
                out[i,j,k,2] = dvx_dz - dvz_dx
                
                # z-component: ∂vᵧ/∂x - ∂vₓ/∂y
                dvy_dx = (vector_field[i+1,j,k,2] - vector_field[i-1,j,k,2]) * hx_inv
                dvx_dy = (vector_field[i,j+1,k,1] - vector_field[i,j-1,k,1]) * hy_inv
                out[i,j,k,3] = dvy_dx - dvx_dy
            end
            return out
        end
    else
        error("Curl not implemented for $(ndims)D")
    end
end

"""
    make_vector_laplace(grid::CartesianGrid) -> OperatorType

Create vector Laplacian operator: ∇²v = (∇²vₓ, ∇²vᵧ, ∇²vᵤ)
"""
function make_vector_laplace(grid::CartesianGrid)
    scalar_laplace = make_laplace(grid)
    ndims = ndim(grid)
    
    function vector_laplace(vector_field, out)
        for component in 1:ndims
            # Apply scalar Laplacian to each component
            temp_in = view(vector_field, :, :, component)
            temp_out = view(out, :, :, component)
            scalar_laplace(temp_in, temp_out)
        end
        return out
    end
end

"""
    make_biharmonic(grid::CartesianGrid) -> OperatorType

Create biharmonic operator: ∇⁴ = ∇²(∇²)
"""
function make_biharmonic(grid::CartesianGrid)
    laplace_op = make_laplace(grid)
    
    function biharmonic(field_data, out)
        # Apply Laplacian twice
        temp = similar(field_data)
        laplace_op(field_data, temp)
        laplace_op(temp, out)
        return out
    end
end

"""
    fft_laplace(grid::CartesianGrid) -> OperatorType

Create FFT-based Laplacian for periodic boundary conditions.
"""
function fft_laplace(grid::CartesianGrid)
    dims = size(grid)
    extent_vals = extent(grid)
    ndims = ndim(grid)
    
    # Pre-compute k² values
    k_squared = zeros(dims)
    
    if ndims == 1
        L = extent_vals[1][2] - extent_vals[1][1]
        for i in 1:dims[1]
            k = 2π * (i - 1 - dims[1]÷2) / L
            k_squared[i] = k^2
        end
    elseif ndims == 2
        Lx = extent_vals[1][2] - extent_vals[1][1]
        Ly = extent_vals[2][2] - extent_vals[2][1]
        
        for i in 1:dims[1], j in 1:dims[2]
            kx = 2π * (i - 1 - dims[1]÷2) / Lx
            ky = 2π * (j - 1 - dims[2]÷2) / Ly
            k_squared[i,j] = kx^2 + ky^2
        end
    else
        error("FFT Laplacian not implemented for $(ndims)D")
    end
    
    function fft_laplacian(field_data, out)
        # Forward FFT
        field_fft = fft(field_data)
        
        # Apply -k² in Fourier space
        field_fft .*= -k_squared
        
        # Inverse FFT
        out .= real.(ifft(field_fft))
        return out
    end
end

"""
    make_corner_point_setter_2d(grid::CartesianGrid) -> Function

Create function to set corner points for 2D grid boundary conditions.
"""
function make_corner_point_setter_2d(grid::CartesianGrid)
    # Check for periodicity
    bcs = get_boundary_conditions(grid)
    periodic_x = isa(bcs[1], PeriodicBC)
    periodic_y = isa(bcs[2], PeriodicBC)
    
    function set_corner_points!(arr)
        nx, ny = size(arr)
        
        if periodic_x
            # Use periodicity in x-direction
            arr[1, 1] = arr[end-1, 1]
            arr[end, 1] = arr[2, 1]
            arr[1, end] = arr[end-1, end]
            arr[end, end] = arr[2, end]
        elseif periodic_y
            # Use periodicity in y-direction  
            arr[1, 1] = arr[1, end-1]
            arr[end, 1] = arr[end, 2]
            arr[1, end] = arr[1, 2]
            arr[end, end] = arr[end, end-1]
        else
            # Interpolate corner points
            arr[1, 1] = 0.5 * (arr[1, 2] + arr[2, 1])
            arr[end, 1] = 0.5 * (arr[end, 2] + arr[end-1, 1])
            arr[1, end] = 0.5 * (arr[1, end-1] + arr[2, end])
            arr[end, end] = 0.5 * (arr[end, end-1] + arr[end-1, end])
        end
        
        return arr
    end
end

"""
    make_poisson_solver(grid::CartesianGrid, bcs) -> Function

Create Poisson equation solver: ∇²u = f
"""
function make_poisson_solver(grid::CartesianGrid, bcs)
    return make_general_poisson_solver(grid, bcs)
end